// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <cstdint>
#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::functions::test;

namespace {
class ArrayRemoveNullsTest : public FunctionBaseTest {
 protected:
  void testArrayRemoveNull(const VectorPtr& expected, const VectorPtr& input) {
    auto result = evaluate("remove_nulls(c0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }
};

TEST_F(ArrayRemoveNullsTest, simpleString) {
  auto input = make_nullable_array_vector<std::string>(
      {{"foo", std::nullopt, "bar"}, {"foo", "bar"}});
  auto expected =
      make_nullable_array_vector<std::string>({{"foo", "bar"}, {"foo", "bar"}});
  testArrayRemoveNull(expected, input);
}

TEST_F(ArrayRemoveNullsTest, simpleInt) {
  auto input = make_nullable_array_vector<int>(
      {{1, std::nullopt, std::nullopt, 3}, {1, 3}});
  auto expected = make_nullable_array_vector<int>({{1, 3}, {1, 3}});
  testArrayRemoveNull(expected, input);
}

TEST_F(ArrayRemoveNullsTest, simpleBool) {
  auto input = make_nullable_array_vector<bool>(
      {{true, false, true}, {true, false, std::nullopt}});
  auto expected =
      make_nullable_array_vector<bool>({{true, false, true}, {true, false}});
  testArrayRemoveNull(expected, input);
}

TEST_F(ArrayRemoveNullsTest, complexType) {
  using array_type = std::optional<std::vector<std::optional<int32_t>>>;
  array_type array1 = {{1, 2}};
  array_type array2 = std::nullopt;
  array_type array3 = {{1, std::nullopt, 2}};

  auto input =
      make_nullable_nested_array_vector<int32_t>({{{array1, array2, array3}}});

  array_type earray1 = {{1, 2}};
  array_type earray2 = {{1, std::nullopt, 2}};
  auto expected =
      make_nullable_nested_array_vector<int32_t>({{{earray1, earray2}}});
  testArrayRemoveNull(expected, input);
}

} // namespace
